package main

import (
	"compress/gzip" // 新增：用于解压
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
)

const uploadDir = "./uploads"

func withCORS(h http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}
		h(w, r)
	}
}

func main() {
	http.HandleFunc("/upload", withCORS(uploadHandler))
	http.HandleFunc("/merge", withCORS(mergeHandler))
	http.HandleFunc("/uploaded", withCORS(uploadedHandler))
	http.Handle("/", http.FileServer(http.Dir("./static")))

	fmt.Println("服务已启动: http://localhost:8080")
	http.ListenAndServe(":8080", nil)
}

func uploadedHandler(w http.ResponseWriter, r *http.Request) {
	fileName := r.URL.Query().Get("filename")
	chunkDir := filepath.Join(uploadDir, fileName+".chunks")

	files, err := os.ReadDir(chunkDir)
	if err != nil {
		w.Write([]byte("[]"))
		return
	}

	var uploaded []string
	for _, f := range files {
		if !f.IsDir() {
			uploaded = append(uploaded, f.Name())
		}
	}

	jsonData, _ := json.Marshal(uploaded)
	w.Header().Set("Content-Type", "application/json")
	w.Write(jsonData)
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "只支持 POST", http.StatusMethodNotAllowed)
		return
	}

	chunkIndex := r.FormValue("index")
	fileName := r.FormValue("filename")
	isCompressed := r.FormValue("compressed") == "1" // 检查是否压缩

	file, _, err := r.FormFile("file")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer file.Close()

	var reader io.Reader = file
	if isCompressed {
		gzReader, err := gzip.NewReader(file)
		if err != nil {
			http.Error(w, "解压失败: "+err.Error(), http.StatusInternalServerError)
			return
		}
		defer gzReader.Close()
		reader = gzReader
	}

	chunkDir := filepath.Join(uploadDir, fileName+".chunks")
	os.MkdirAll(chunkDir, os.ModePerm)

	chunkPath := filepath.Join(chunkDir, chunkIndex)
	out, err := os.Create(chunkPath)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer out.Close()

	_, err = io.Copy(out, reader)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Write([]byte("分片上传成功: " + chunkIndex))
}

func mergeHandler(w http.ResponseWriter, r *http.Request) {
	fileName := r.URL.Query().Get("filename")
	totalChunksStr := r.URL.Query().Get("total")
	totalChunks, err := strconv.Atoi(totalChunksStr)
	if err != nil || totalChunks <= 0 {
		http.Error(w, "无效的 total 参数", http.StatusBadRequest)
		return
	}

	chunkDir := filepath.Join(uploadDir, fileName+".chunks")
	finalFilePath := filepath.Join(uploadDir, fileName)

	// 获取第一个分片大小（假设除了最后一片都相同）
	firstChunkPath := filepath.Join(chunkDir, "0")
	firstStat, err := os.Stat(firstChunkPath)
	if err != nil {
		http.Error(w, "找不到分片 0: "+err.Error(), http.StatusInternalServerError)
		return
	}
	chunkSize := firstStat.Size()

	// 创建最终文件，如果已存在则删除
	if stat, err := os.Stat(finalFilePath); err == nil {
		if stat.IsDir() {
			os.RemoveAll(finalFilePath)
		} else {
			os.Remove(finalFilePath)
		}
	}
	finalFile, err := os.Create(finalFilePath)
	if err != nil {
		http.Error(w, "创建合并文件失败: "+err.Error(), http.StatusInternalServerError)
		return
	}
	defer finalFile.Close()

	// 预分配文件大小
	err = finalFile.Truncate(chunkSize * int64(totalChunks))
	if err != nil {
		http.Error(w, "预分配文件大小失败: "+err.Error(), http.StatusInternalServerError)
		return
	}

	// 并行写入
	errChan := make(chan error, totalChunks)

	for i := 0; i < totalChunks; i++ {
		i := i // 避免闭包变量陷阱
		go func() {
			chunkPath := filepath.Join(chunkDir, strconv.Itoa(i))
			data, err := os.ReadFile(chunkPath)
			if err != nil {
				errChan <- fmt.Errorf("读取分片 %d 失败: %w", i, err)
				return
			}

			// 写入对应位置
			_, err = finalFile.WriteAt(data, int64(i)*chunkSize)
			if err != nil {
				errChan <- fmt.Errorf("写入分片 %d 失败: %w", i, err)
				return
			}

			// 删除分片
			os.Remove(chunkPath)
			errChan <- nil
		}()
	}

	// 等待所有分片写入完成
	for i := 0; i < totalChunks; i++ {
		if err := <-errChan; err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}

	// 删除分片目录
	os.RemoveAll(chunkDir)

	w.Write([]byte("文件合并完成: " + finalFilePath))
}
